package resource

import (
	"fmt"
	"sort"
	"sync"

	"github.com/louloulin/dataflare/pkg/controller/membership"
)

// ResourceType 表示资源类型
type ResourceType string

const (
	// ResourceTypeCPU 表示CPU资源
	ResourceTypeCPU ResourceType = "cpu"
	// ResourceTypeMemory 表示内存资源
	ResourceTypeMemory ResourceType = "memory"
	// ResourceTypeDisk 表示磁盘资源
	ResourceTypeDisk ResourceType = "disk"
	// ResourceTypeNetwork 表示网络资源
	ResourceTypeNetwork ResourceType = "network"
)

// ResourceQuantity 表示资源数量
type ResourceQuantity float64

// ResourceRequirement 表示资源需求
type ResourceRequirement struct {
	// Type 是资源类型
	Type ResourceType
	// Quantity 是资源数量
	Quantity ResourceQuantity
}

// ResourceCapacity 表示资源容量
type ResourceCapacity struct {
	// Type 是资源类型
	Type ResourceType
	// Total 是总容量
	Total ResourceQuantity
	// Used 是已使用容量
	Used ResourceQuantity
}

// Available 返回可用容量
func (rc *ResourceCapacity) Available() ResourceQuantity {
	return rc.Total - rc.Used
}

// UtilizationRatio 返回利用率
func (rc *ResourceCapacity) UtilizationRatio() float64 {
	if rc.Total == 0 {
		return 0
	}
	return float64(rc.Used) / float64(rc.Total)
}

// NodeResources 表示节点资源
type NodeResources struct {
	// NodeID 是节点ID
	NodeID string
	// Capacities 是资源容量映射表
	Capacities map[ResourceType]*ResourceCapacity
	// mu 是互斥锁
	mu sync.RWMutex
}

// NewNodeResources 创建一个新的节点资源
func NewNodeResources(nodeID string) *NodeResources {
	return &NodeResources{
		NodeID:     nodeID,
		Capacities: make(map[ResourceType]*ResourceCapacity),
	}
}

// SetCapacity 设置资源容量
func (nr *NodeResources) SetCapacity(resourceType ResourceType, total ResourceQuantity) {
	nr.mu.Lock()
	defer nr.mu.Unlock()

	if capacity, ok := nr.Capacities[resourceType]; ok {
		capacity.Total = total
	} else {
		nr.Capacities[resourceType] = &ResourceCapacity{
			Type:  resourceType,
			Total: total,
			Used:  0,
		}
	}
}

// AllocateResource 分配资源
func (nr *NodeResources) AllocateResource(resourceType ResourceType, quantity ResourceQuantity) bool {
	nr.mu.Lock()
	defer nr.mu.Unlock()

	capacity, ok := nr.Capacities[resourceType]
	if !ok {
		return false
	}

	if capacity.Available() < quantity {
		return false
	}

	capacity.Used += quantity
	return true
}

// ReleaseResource 释放资源
func (nr *NodeResources) ReleaseResource(resourceType ResourceType, quantity ResourceQuantity) bool {
	nr.mu.Lock()
	defer nr.mu.Unlock()

	capacity, ok := nr.Capacities[resourceType]
	if !ok {
		return false
	}

	if capacity.Used < quantity {
		capacity.Used = 0
	} else {
		capacity.Used -= quantity
	}
	return true
}

// GetCapacity 获取资源容量
func (nr *NodeResources) GetCapacity(resourceType ResourceType) (*ResourceCapacity, bool) {
	nr.mu.RLock()
	defer nr.mu.RUnlock()

	capacity, ok := nr.Capacities[resourceType]
	return capacity, ok
}

// GetAllCapacities 获取所有资源容量
func (nr *NodeResources) GetAllCapacities() map[ResourceType]*ResourceCapacity {
	nr.mu.RLock()
	defer nr.mu.RUnlock()

	capacities := make(map[ResourceType]*ResourceCapacity, len(nr.Capacities))
	for resourceType, capacity := range nr.Capacities {
		capacities[resourceType] = &ResourceCapacity{
			Type:  capacity.Type,
			Total: capacity.Total,
			Used:  capacity.Used,
		}
	}
	return capacities
}

// ResourceAllocator 表示资源分配器
type ResourceAllocator struct {
	// nodeResources 是节点资源映射表
	nodeResources map[string]*NodeResources
	// mu 是互斥锁
	mu sync.RWMutex
}

// NewResourceAllocator 创建一个新的资源分配器
func NewResourceAllocator() *ResourceAllocator {
	return &ResourceAllocator{
		nodeResources: make(map[string]*NodeResources),
	}
}

// RegisterNode 注册节点
func (ra *ResourceAllocator) RegisterNode(nodeID string) *NodeResources {
	ra.mu.Lock()
	defer ra.mu.Unlock()

	if resources, ok := ra.nodeResources[nodeID]; ok {
		return resources
	}

	resources := NewNodeResources(nodeID)
	ra.nodeResources[nodeID] = resources
	return resources
}

// UnregisterNode 注销节点
func (ra *ResourceAllocator) UnregisterNode(nodeID string) bool {
	ra.mu.Lock()
	defer ra.mu.Unlock()

	if _, ok := ra.nodeResources[nodeID]; ok {
		delete(ra.nodeResources, nodeID)
		return true
	}
	return false
}

// GetNodeResources 获取节点资源
func (ra *ResourceAllocator) GetNodeResources(nodeID string) (*NodeResources, bool) {
	ra.mu.RLock()
	defer ra.mu.RUnlock()

	resources, ok := ra.nodeResources[nodeID]
	return resources, ok
}

// GetAllNodeResources 获取所有节点资源
func (ra *ResourceAllocator) GetAllNodeResources() map[string]*NodeResources {
	ra.mu.RLock()
	defer ra.mu.RUnlock()

	nodeResources := make(map[string]*NodeResources, len(ra.nodeResources))
	for nodeID, resources := range ra.nodeResources {
		nodeResources[nodeID] = resources
	}
	return nodeResources
}

// AllocateTask 分配任务
func (ra *ResourceAllocator) AllocateTask(requirements []ResourceRequirement) (string, error) {
	ra.mu.RLock()
	defer ra.mu.RUnlock()

	if len(ra.nodeResources) == 0 {
		return "", fmt.Errorf("no nodes available")
	}

	// 找到最适合的节点
	bestNodeID := ra.findBestNode(requirements)
	if bestNodeID == "" {
		return "", fmt.Errorf("no suitable node found")
	}

	// 分配资源
	nodeResources := ra.nodeResources[bestNodeID]
	for _, requirement := range requirements {
		if !nodeResources.AllocateResource(requirement.Type, requirement.Quantity) {
			// 回滚已分配的资源
			for _, req := range requirements {
				if req.Type == requirement.Type {
					break
				}
				nodeResources.ReleaseResource(req.Type, req.Quantity)
			}
			return "", fmt.Errorf("failed to allocate %v resource", requirement.Type)
		}
	}

	return bestNodeID, nil
}

// ReleaseTask 释放任务
func (ra *ResourceAllocator) ReleaseTask(nodeID string, requirements []ResourceRequirement) error {
	ra.mu.RLock()
	defer ra.mu.RUnlock()

	nodeResources, ok := ra.nodeResources[nodeID]
	if !ok {
		return fmt.Errorf("node %s not found", nodeID)
	}

	for _, requirement := range requirements {
		if !nodeResources.ReleaseResource(requirement.Type, requirement.Quantity) {
			return fmt.Errorf("failed to release %v resource", requirement.Type)
		}
	}

	return nil
}

// findBestNode 找到最适合的节点
func (ra *ResourceAllocator) findBestNode(requirements []ResourceRequirement) string {
	// 检查每个节点是否满足要求
	candidates := make([]string, 0, len(ra.nodeResources))
	for nodeID, resources := range ra.nodeResources {
		suitable := true
		for _, requirement := range requirements {
			capacity, ok := resources.GetCapacity(requirement.Type)
			if !ok || capacity.Available() < requirement.Quantity {
				suitable = false
				break
			}
		}
		if suitable {
			candidates = append(candidates, nodeID)
		}
	}

	if len(candidates) == 0 {
		return ""
	}

	// 如果只有一个候选节点，直接返回
	if len(candidates) == 1 {
		return candidates[0]
	}

	// 根据资源利用率排序
	// 这里使用最低利用率的节点，以实现负载均衡
	type nodeScore struct {
		nodeID string
		score  float64
	}
	scores := make([]nodeScore, 0, len(candidates))

	for _, nodeID := range candidates {
		resources := ra.nodeResources[nodeID]
		// 计算平均利用率
		var totalRatio float64
		var count int
		for _, requirement := range requirements {
			if capacity, ok := resources.GetCapacity(requirement.Type); ok {
				totalRatio += capacity.UtilizationRatio()
				count++
			}
		}
		avgRatio := totalRatio / float64(count)
		scores = append(scores, nodeScore{nodeID: nodeID, score: avgRatio})
	}

	// 按分数排序
	sort.Slice(scores, func(i, j int) bool {
		return scores[i].score < scores[j].score
	})

	return scores[0].nodeID
}

// ResourceAllocationManager 表示资源分配管理器
type ResourceAllocationManager struct {
	// allocator 是资源分配器
	allocator *ResourceAllocator
	// clusterManager 是集群管理器
	clusterManager *membership.ClusterManager
	// mu 是互斥锁
	mu sync.RWMutex
}

// NewResourceAllocationManager 创建一个新的资源分配管理器
func NewResourceAllocationManager(clusterManager *membership.ClusterManager) *ResourceAllocationManager {
	manager := &ResourceAllocationManager{
		allocator:      NewResourceAllocator(),
		clusterManager: clusterManager,
	}

	// 注册集群事件处理
	clusterManager.AddDelegate(manager)

	return manager
}

// OnMemberJoin 处理成员加入事件
func (ram *ResourceAllocationManager) OnMemberJoin(member *membership.Member) {
	ram.allocator.RegisterNode(member.ID)
}

// OnMemberLeave 处理成员离开事件
func (ram *ResourceAllocationManager) OnMemberLeave(member *membership.Member) {
	ram.allocator.UnregisterNode(member.ID)
}

// OnMemberFailed 处理成员失败事件
func (ram *ResourceAllocationManager) OnMemberFailed(member *membership.Member) {
	ram.allocator.UnregisterNode(member.ID)
}

// OnMemberUpdate 处理成员更新事件
func (ram *ResourceAllocationManager) OnMemberUpdate(member *membership.Member) {
	// 这里可以处理成员资源更新
}

// AllocateTask 分配任务
func (ram *ResourceAllocationManager) AllocateTask(requirements []ResourceRequirement) (string, error) {
	return ram.allocator.AllocateTask(requirements)
}

// ReleaseTask 释放任务
func (ram *ResourceAllocationManager) ReleaseTask(nodeID string, requirements []ResourceRequirement) error {
	return ram.allocator.ReleaseTask(nodeID, requirements)
}

// GetNodeResources 获取节点资源
func (ram *ResourceAllocationManager) GetNodeResources(nodeID string) (*NodeResources, bool) {
	return ram.allocator.GetNodeResources(nodeID)
}

// GetAllNodeResources 获取所有节点资源
func (ram *ResourceAllocationManager) GetAllNodeResources() map[string]*NodeResources {
	return ram.allocator.GetAllNodeResources()
}

// GetAllocator 获取资源分配器
func (ram *ResourceAllocationManager) GetAllocator() *ResourceAllocator {
	return ram.allocator
}

// UpdateNodeCapacity 更新节点容量
func (ram *ResourceAllocationManager) UpdateNodeCapacity(nodeID string, resourceType ResourceType, total ResourceQuantity) error {
	resources, ok := ram.allocator.GetNodeResources(nodeID)
	if !ok {
		return fmt.Errorf("node %s not found", nodeID)
	}

	resources.SetCapacity(resourceType, total)
	return nil
}
